# new_predict.py
import torch
import joblib
import numpy as np
import json
from models.price_predictor import PricePredictor

NUMERIC_COLUMNS = ["预算", "流行度", "平均评分", "评分人数", "已上映天数", "入座率", "影院评分"]


def apply_adjustments(base_price, city_level=None, session=None, user_tags=None, movie_tags=None):
    """根据城市、场次、偏好调整预测票价，最大±10元"""
    adjustment = 0.0

    # 城市等级调整（高等级城市价格更低）
    if city_level is not None:
        city_adjust = {1: 0.0, 2: -2, 3: -4, 4: -6, 5: -8}
        adjustment += city_adjust.get(city_level, 0.0)

    # 场次时间调整（越晚越贵）
    if session is not None:
        session_adjust = {"morning": -3, "noon": 0, "afternoon": 4, "evening": 3}
        adjustment += session_adjust.get(session, 0.0)

    # 偏好重合度调整（每重合1项 +1 元）
    if user_tags and movie_tags:
        overlap = len(set(user_tags) & set(movie_tags))
        adjustment += overlap * 1.0

    # 控制总调整幅度在 [-10, 10]
    adjustment = max(min(adjustment, 10), -10)
    return base_price + adjustment

def predict_price_from_features(feature_dict, city_level=None, session=None,
                                user_tags=None, movie_tags=None,
                                model_path="movie-ticket-bidding/checkpoints/ResidualBlockGELU/best_model.pt",
                                config_path="movie-ticket-bidding/config/ResidualBlockGELU.json",
                                scaler_path="movie-ticket-bidding/data/processed/scaler.pkl"):
    for key in NUMERIC_COLUMNS:
        if key not in feature_dict:
            raise ValueError(f"缺少特征: {key}")

    input_vals = np.array([feature_dict[k] for k in NUMERIC_COLUMNS]).reshape(1, -1)
    scaler = joblib.load(scaler_path)
    input_scaled = scaler.transform(input_vals)
    x = torch.tensor(input_scaled, dtype=torch.float32)

    with open(config_path, "r") as f:
        model_cfg = json.load(f)["model"]

    input_dim = x.shape[1]
    model = PricePredictor(config_path, input_dim)
    model.load_state_dict(torch.load(model_path, weights_only=True))
    model.eval()

    with torch.no_grad():
        output = model(x)
        base_price = output.item()

    # 根据规则调整价格
    final_price = apply_adjustments(base_price, city_level, session, user_tags, movie_tags)
    return final_price


if __name__ == "__main__":
    example = {
        "预算": 237000000,
        "流行度": 169.4,
        "平均评分": 8.4,
        "评分人数": 218597,
        "已上映天数": 30,
        "入座率": 0,
        "影院评分": 5
    }
    result = predict_price_from_features(
        feature_dict=example,
        city_level=3,
        session="evening",
        user_tags=["动作", "悬疑"],
        movie_tags=["动作", "战争", "科幻"]
    )
    if result < 25:
        result = 25
    print(f" 预测票价: {result:.2f} 元")